package zk.tool;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import zk.tool.model.TreeItemData;
import zk.tool.util.StringUtil;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.function.Function;

public class ZKClient implements Watcher {
    private ZooKeeper zooKeeper;
    private int timeout = 4000;
    private String connectString;
    private boolean auth = false;
    private String user;
    private String password;
    private volatile boolean connectSuccess = false;
    private CountDownLatch countDownLatch = new CountDownLatch(1);

    public ZKClient(String connectString) {
        this.connectString = connectString;
    }

    @Override
    public void process(WatchedEvent watchedEvent) {
        System.out.println(watchedEvent);
        if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
            System.out.println("Watch received event");
            this.connectSuccess = true;
            countDownLatch.countDown();
        }
    }

    public void connect() throws IOException, InterruptedException {
        this.zooKeeper = new ZooKeeper(this.connectString, this.timeout, this);
        if (StringUtil.isNotBlank(user) && StringUtil.isNotBlank(password)) {
            this.zooKeeper.addAuthInfo(user, this.password.getBytes(StandardCharsets.UTF_8));
        }
        this.countDownLatch.await();
        System.out.println("zookeeper connection success");
    }

    public String createNode(String path, String data) throws InterruptedException, KeeperException {
        if (this.isAuth()) {
            return this.zooKeeper.create(
                    path,
                    data.getBytes(StandardCharsets.UTF_8),
                    ZooDefs.Ids.CREATOR_ALL_ACL,
                    CreateMode.PERSISTENT
            );
        }
        return this.zooKeeper.create(
                path,
                data.getBytes(StandardCharsets.UTF_8),
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.PERSISTENT
        );
    }

    public List<String> getChildren(String path) throws InterruptedException, KeeperException {
        Stat exists = this.zooKeeper.exists(path, false);
        if (exists != null) {
            return zooKeeper.getChildren(path, false);
        }
        return Collections.emptyList();
    }

    public TreeItemData getData(String path) throws InterruptedException, KeeperException {
        TreeItemData itemData = new TreeItemData("", "");
        Stat stat = new Stat();
        byte[] data = zooKeeper.getData(path, false, stat);
        if (data == null) {
            return itemData;
        }
        itemData.setData(new String(data, StandardCharsets.UTF_8));
        itemData.setStat(stat);
        return itemData;
    }

    public Stat setData(String path, String data) throws InterruptedException, KeeperException {
        return zooKeeper.setData(path, data.getBytes(StandardCharsets.UTF_8), -1);
    }

    public void delete(String path) throws InterruptedException, KeeperException {
        zooKeeper.delete(path, -1);
    }

    public boolean isExists(String path) throws InterruptedException, KeeperException {
        return zooKeeper.exists(path, false) != null;
    }

    public void close() {
        if (this.zooKeeper != null) {
            try {
                this.zooKeeper.close();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    protected void finalize() throws Throwable {
        super.finalize();
        if (this.zooKeeper != null) {
            try {
                this.zooKeeper.close();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public ZooKeeper getZooKeeper() {
        return zooKeeper;
    }

    public String getConnectString() {
        return connectString;
    }

    public boolean isConnectSuccess() {
        return connectSuccess;
    }

    public int getTimeout() {
        return timeout;
    }

    public void setTimeout(int timeout) {
        this.timeout = timeout;
    }

    public String getUser() {
        return user;
    }

    public void setUser(String user) {
        this.user = user;
    }

    public String getPassword() {
        return password;
    }

    public boolean isAuth() {
        return auth;
    }

    public void setAuth(boolean auth) {
        this.auth = auth;
    }

    public void setPassword(String password) {
        this.password = password;
    }
}
